
#ifndef _MATRIXSOLVER_HPP
#define _MATRIXSOLVER_HPP

#include "Matrix.hpp"
#include <algorithm>
#include <limits>
#include <cmath>

// Solvers for Matrix class
// A square matrix is required
// $Id: MatrixSolver.hpp 2 2012-05-22 17:55:16Z gerry@gaboury.biz $

template <class T> class MatrixSolver : public Matrix<T>
{

private:
  // # pivot vectors, lower & upper bands
  int *m_piv,*m_ipiv,*m_jpiv;
  bool m_isDecomp;
  bool m_isInvert;

public:

  // constructors

  MatrixSolver () : 
    Matrix<T>(), m_piv(0), m_ipiv(0), m_jpiv(0),
    m_isDecomp(false), m_isInvert(false) { }

  MatrixSolver ( int p_row, int p_col ) : 
    Matrix<T>(p_row,p_col), m_piv(0), m_ipiv(0), m_jpiv(0),
    m_isDecomp(false), m_isInvert(false) { }

  MatrixSolver ( int p_row, int p_col, T p_val ) : 
    Matrix<T>(p_row,p_col,p_val), m_piv(0), m_ipiv(0), m_jpiv(0),
    m_isDecomp(false), m_isInvert(false) { }

  MatrixSolver ( const Matrix<T> & p_mat ) :
    Matrix<T> (p_mat),
    m_piv(0), m_ipiv(0), m_jpiv(0),
    m_isDecomp(false), m_isInvert(false) { }

  virtual ~MatrixSolver() { delete [] m_piv,m_ipiv,m_jpiv; }

  // decomposition and solvers
  // LU decomposition - partial pivoting
  virtual void  LU_Decomp ( const T& p_eps=0 )
  {
    T eps = max ( numeric_limits<T>::epsilon(), p_eps );

    // arbitrary precision has an eps of zero so use
    // long double as the lower limit
    double ldeps = numeric_limits<long double>::epsilon();

    if ( eps < ldeps )
      eps = ldeps;
    
    if ( this->row() != this->col() )
      throw (range_error("MatrixSolver::LU : Matrix not square"));

    int j,k,irow,jmax,n(this->row());
    T r,d,amax,zero(0);

    // matrix already decomposed
    if ( m_isDecomp )
      throw (logic_error("MatrixSolver::LU : matrix already decomposed"));

//     if ( m_lower >= 0 && m_upper >= 0 )
//       {
// 	try
// 	  {
// 	    this->LU_Banded();
// 	  }
// 	catch (Error& err)
// 	  {
// 	    throw;
// 	  }
// 	return;
//       }

    if ( !m_piv )
      m_piv = new int[n];

    /* Initialize pivoting vector */
    for (k = 0; k < n; ++k)
      m_piv[k] = k;

    for (j = 0; j < n-1; ++j) // j < n-1
      { 
	Vector<T>& Aj = (*this)[j]; // jth row
	/* Partial pivoting */
	jmax = j;
	amax = std::abs( (*this)[j][j] );
	for (irow = j+1; irow < n; ++irow)
	  if (amax < std::abs((*this)[irow][j]))
	    {
	      jmax = irow;
	      amax = std::abs((*this)[irow][j]);
	    }
	if (jmax != j)
	  {
	    std::swap ( m_piv[j], m_piv[jmax] );
	    Vector<T>& Ar = (*this)[jmax];
	    for (k=0; k < n; ++k)
	      std::swap ( Aj[k], Ar[k] );
	  }
	/* GAUSSIAN ELIMINATION */
	d = Aj[j];
	if ( std::abs(d) < eps )
	  throw (underflow_error("MatrixSolver::LU : singular matrix"));

	for (irow = j+1; irow < n; ++irow)
	  {
	    Vector<T>& Ar = (*this)[irow];
	    r = Ar[j];
	    if (r == zero) continue;
	    r /= d;
	    Ar[j] = r;
	    for (k=j+1; k < n; ++k)
	      Ar[k] -= r * Aj[k];
	  }
      }
    m_isDecomp = true;	// indicates a decomposed matrix
  }
  //void    LU_fp();		// full pivot
  //void    LU_Banded();		// Banded matrices
  //Vector& Improve ( Matrix&, Vector&, Vector& );
  // Solver
  virtual Vector<T>  LU_Solve ( const Vector<T>& p_vec )
  {
    int k,j;
    int n(this->row());
    T tmp;

    if ( n != p_vec.size() )
      throw (range_error("MatrixSolver::LUSolve : dimension error"));

    // check for decomposition
    if ( !m_isDecomp )
      throw (logic_error("MatrixSolver::LUSolve : matrix not decomposed"));

    Vector<T> Xvec(n);

    /* pivot X */
    for (k=0; k < n; ++k)
      Xvec[k] = p_vec[m_piv[k]];

    /* solve Ly = b */
    for (j = 1; j < n; ++j)
      {
	Vector<T>& Aj = (*this)[j];
	tmp = Xvec[j];
	for (k=0; k < j; ++k) 
	  tmp -= Xvec[k] * Aj[k];
	Xvec[j] = tmp;
      }

    /*  solve Ux = y */

    j = n-1;
    Xvec[j] /= (*this)[j][j];
    for (j=n-2; j >= 0; --j)
      {
	Vector<T>& Aj = (*this)[j];
	tmp = Xvec[j];
	for (k=j+1; k < n; ++k)
	  tmp -= Aj[k] * Xvec[k];
	tmp /= Aj[j];
	Xvec[j] = tmp;
      }
    return Xvec;
  }
  //Vector  LUSolve_fp ( Vector& );
  //Vector  LUSolve_Banded ( Vector& );	// Banded Solver
  // determinant
  T Det()
  { 
    // check for decomposition
    if ( m_isInvert )
      throw (logic_error("MatrixSolver::Det : matrix inverted"));

    if ( !m_isDecomp )
      throw (logic_error("MatrixSolver::Det : matrix not decomposed"));

    T det=1;
    for ( int k=0; k < this->row(); ++k )
      det *= (*this)[k][k];
    return (det);
  }

  // In-place Inversion
  void Invert()
  {
    int i,j,k,n=this->row();
    double dtmp,r,d;

    // check for decomposition
    if ( ! m_isDecomp )
      throw (logic_error("MatrixSolver::Invert : matrix not decomposed"));

    if ( m_isInvert )
      throw (logic_error("MatrixSolver::Invert : matrix inverted"));

    // invert L
    //Vector<T>& Ai = (*this)[1];
    (*this)[1][0] = -(*this)[1][0];
    for (i=2; i < n; ++i)
      {
	Vector<T>& Ai = (*this)[i];
	for (j=0; j < i-1; ++j)
	  {
	    dtmp = -Ai[j];
	    for (k=j+1;k < i; ++k)
	      dtmp -= Ai[k] * (*this)[k][j];
	    Ai[j] = dtmp;
	  }
	Ai[i-1] = -Ai[i-1];
      }

    //  solve U(A**-1) = (L**-1)
    
    Vector<T>& A0 = (*this)[0];
    d = 1.0 / A0[0];
    A0[0]= d;
    for (k=1; k < n; ++k)
      A0[k] *= d;
    for (i=1; i < n; ++i)
      {
	Vector<T>& Ai = (*this)[i];
	d = 1.0 / Ai[i];
	Ai[i]= 1.0;
	for (k=0; k < n; ++k)
	  Ai[k] *= d;
	for (j=0; j < i; ++j)
	  {
	    Vector<T>& Aj = (*this)[j];
	    r = Aj[i];
	    Aj[i] = 0.0;
	    for (k=0; k < n; ++k)
	      Aj[k] -= r * Ai[k];
	  }
      }

    // reverse pivot
    for (j=0; j < n-1; ++j)
      {
	for (k=j; k < n; ++k)
	  if (j == m_piv[k])
	    break;
	if (k == j)
	  continue;
	m_piv[k] = m_piv[j];
	for (i=0; i < n; ++i)
	  {
	    Vector<T>& Ai = (*this)[i];
	    std::swap (Ai[j],Ai[k]); // exchange columns j and k
	  }
      }
    m_isInvert = true;
  }

  // Cholesky decomposition of a symmetric positive definite matrix.
  // Golub & Van Loan p.89
  // A = GG* where G is lower triangular
  //
  void Cholesky_Decomp()
  {
    if ( ! this->isSymmetric() )
      throw (logic_error("Cholsky_Decomp: matrix is not symmetric."));

    int n = this->size();

    for ( int k = 0; k < n; ++k )
      {
	T akk = (*this)[k][k];
	for ( int p = 0; p <= k-1; ++p )
	  {
	    T akp = (*this)[k][p];
	    akk -= akp * akp;
	  }
	akk = sqrt ( akk );
	(*this)[k][k] = akk;
	for ( int i=k+1; i < n; ++i )
	  {
	    T aik = (*this)[i][k];
	    for ( int p = 0; p <= k-1; ++p )
	      aik -= (*this)[i][p] * (*this)[k][p];

	    (*this)[i][k] = aik / akk;
	  }
      }
    m_isDecomp = true;
  }
  Vector<T> Cholesky_Solve( const Vector<T>& p_vec )
  {
    if ( ! m_isDecomp )
      this->Cholesky_Decomp();
    //throw (logic_error("Cholsky_Solve: matrix not decomposed"));

    int n = this->size();

    if ( p_vec.size() != n )
      throw (logic_error("Cholsky_Solve: Vector size error"));

    Vector<T> Vtmp ( p_vec );

    // solve Gy = b
    for ( int k = 0; k < n; ++k )
      {
	T yk = Vtmp[k];

	for ( int p = 0; p <= k-1; ++p )
	  yk -= (*this)[k][p] * Vtmp[p];

	Vtmp[k] = yk / (*this)[k][k];
      }
    // solve G^x = y
    for ( int k = n-1; k >=0 ; --k )
      {
	T xk = Vtmp[k];

	for ( int p = n-1; p >= k+1; --p )
	  xk -= (*this)[p][k] * Vtmp[p];

	Vtmp[k] = xk / (*this)[k][k];
      }
    return Vtmp;
  }
  Matrix<T> Cholesky_Inverse()
  {
    if ( ! m_isDecomp )
      this->Cholesky_Decomp();
    //throw (logic_error("Cholsky_Invert: matrix not decomposed."));

    int n = this->size();
    Matrix<T> Mtmp;

    // solve the identity matrix
    for ( int i=0; i < n; ++i )
      {
	Vector<T> Vec(n,0.0);
	Vec[i] = 1.0;
	Mtmp.push_back(this->Cholesky_Solve(Vec));
      }
    return Mtmp;
  }
  //
  // Conjugate gradient for symmetric positive definite matrix.
  //
  Vector<T> Conjugate_Gradient( const Vector<T>& p_vec )
  {
    if ( ! this->isSymmetric() )
      throw (logic_error("Conjugate Gradient: matrix is not symmetric."));

    T eps = numeric_limits<T>::epsilon();

    // arbitrary precision has an eps of zero so use
    // long double as the lower limit
    double ldeps = numeric_limits<long double>::epsilon();

    int n = this->size();

    if ( eps < ldeps )
      eps = ldeps / 4;

    // convergence criterion
    T conv = eps * p_vec.SumSq();

    Vector<T> r ( p_vec );
    Vector<T> p(p_vec.size()), x(p_vec.size(),0.0);
    int k = 1;
    vector<T> rho;
    rho.push_back(r.SumSq());

    T beta;

    do
      {
	if ( k <= 1 )
	  {
	    p = r;
	  }
	else
	  {
	    beta = rho[k-1] / rho[k-2];
	    p = r + p * beta;
	  }
	Vector<T> w = (*this) * p;
	T alpha = rho[k-1] / ( p * w );
	x += p * alpha;
	r -= w * alpha;
	rho.push_back(r.SumSq());
	++k;
      } while ( sqrt(rho[k-1]) > conv && k < 20 );

    //cout << "Conjugate Gradient: " << k << " iterations" << endl;
    return x;
  }

};

#endif	// _MATRIXSOLVER_HPP
